///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/fqmc/orb_opt.cpp
///
/// \brief   implementation file for orbital opitimzation
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<numeric>
#include<cassert>
#include<algorithm>
#include<cmath>
#include<complex>

#include<boost/format.hpp>
#include<boost/shared_ptr.hpp>

#include<formic/exception.h>
#include<formic/archive.h>
#include<formic/all_data.h>
#include<formic/input/base.h>
#include<formic/mpi/interface.h>
#include<formic/numeric/numeric.h>
#include<formic/output/output.h>
#include<formic/lapack/interface.h>
#include<formic/wfn/configuration.h>
#include<formic/wfn/coeff_obj.h>
#include<formic/fqmc/accumulator.h>
#include<formic/vector/vector.h>
#include<formic/vector/matrix.h>
#include<formic/ham/tei.h>
#include<formic/ham/ham.h>
#include<formic/lbfgs.h>
#include<formic/sr1.h>
#include<formic/fqmc/energy_acc.h>
#include<formic/fqmc/gradient_acc.h>
#include<formic/fqmc/sampler.h>
#include<formic/fqmc/orb_opt_acc.h>
#include<formic/fqmc/orb_opt.h>
#include<formic/fqmc/optimizer.h>

namespace formic { namespace fqmc {

  ///////////////////////////////////////////////////////////////////////////////////////////////////
  /// \brief  a class to act as the worker in an lbfgs optimization of the orbitals
  ///
  /// The orbital optimization finds the unitary orbital transformation that minimizaes the wave
  /// function's energy.  This transformation is parameterized by an anti-hermitian matrix, which
  /// upon exponentiation gives the unitary orbital transformation.  The variables that are optimized
  /// are the elements of the anti-hermitian matrix's lower triangle and are referred to here as the
  /// transformation's lower triangle variables.
  ///
  ///////////////////////////////////////////////////////////////////////////////////////////////////
  template<class S> class OrbOptWorker {

    // private data members
    private:

      int _norb; ///< number of orbitals
      int _nlt;  ///< number of lower triangle variables

      bool _couple_bfgs;  ///< whether to include the wave function variables and their gradients in the lbfgs search directions (in addition to those for orbital rotation)
      bool _interstitial_wfn;  ///< whether or not to optimize the wave function variables in between each lbfgs step
      bool _rotate_fixed_pm;  ///< whether or not to rotate the initial fixed pairing matrix at each step so that the AGP function becomes orbital rotation independent
      bool _new_sample_each_iter; ///< whether or not to update the guiding function and take a new MC sample at each BFGS iteration

      formic::Vector<S,int> _lt;     /// lower triangle of the anti-hermitian matrix
      formic::Vector<S,int> _lt_adj; /// derivative adjoints of the lower triangle of the anti-hermitian matrix

      formic::Matrix<S,int> _ahm;     /// anti-hermitian matrix
      formic::Matrix<S,int> _ahm_adj; /// derivative adjoints of the anti-hermitian matrix

      formic::Matrix<S,int> _uot;     /// unitary orbital transform
      formic::Matrix<S,int> _uot_adj; /// derivative adjoints of the unitary orbital transform

      formic::Matrix<S,int> _oei_pq; ///< 0-index-transformed one-electron integrals
      formic::Matrix<S,int> _oei_qi; ///< 1-index-transformed one-electron integrals
      formic::Matrix<S,int> _oei_ij; ///< 2-index-transformed one-electron integrals

      formic::Matrix<S,int> _oei_adj_qi; ///< 1-index-transformed one-electron integral derivative adjoints
      formic::Matrix<S,int> _oei_adj_ij; ///< 2-index-transformed one-electron integral derivative adjoints

      formic::Matrix<S,int> _tei_pqrs; ///< 0-index-transformed two-electron integrals
      formic::Matrix<S,int> _tei_qrsi; ///< 1-index-transformed two-electron integrals
      formic::Matrix<S,int> _tei_rsij; ///< 2-index-transformed two-electron integrals
      formic::Matrix<S,int> _tei_sijk; ///< 3-index-transformed two-electron integrals
      formic::Matrix<S,int> _tei_ijkl; ///< 4-index-transformed two-electron integrals

      formic::Matrix<S,int> _tei_adj_qrsi; ///< 1-index-transformed two-electron integral derivative adjoints
      formic::Matrix<S,int> _tei_adj_rsij; ///< 2-index-transformed two-electron integral derivative adjoints
      formic::Matrix<S,int> _tei_adj_sijk; ///< 3-index-transformed two-electron integral derivative adjoints
      formic::Matrix<S,int> _tei_adj_ijkl; ///< 4-index-transformed two-electron integral derivative adjoints

      formic::Matrix<S,int> _initial_pm; ///< holds the initial pairing matrix if necessary

      boost::shared_ptr<formic::CoeffObj<S> > _co; ///< pointer to the wave function's coefficient object
      boost::shared_ptr<formic::CoeffObj<S> > _old_co; ///< pointer to a saved copy of the coefficient object
      boost::shared_ptr<formic::fqmc::Sampler<S> > _sampler; ///< object to perform the necessary sampling

      formic::Hamiltonian<S> & _ham; ///< reference to the hamiltonian

      const formic::InputBase & _userinp; ///< reference to the user's input options

      std::vector<formic::Vector<S,int> > _pmt_set_history; ///< vector of previously minimized parameter sets
      std::vector<formic::Vector<double,int> > _x_history; ///< vector of previous x vector values

    // private member functions
    private:

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  computes the unitary orbital transforamtion
      ///
      /// \param[in]      n      number of parameters that specify the orbital transformation
      /// \param[in]      x      the elements of the antihermitian matrix's lower triangle, size _nlt
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void build_uot(const int n, const double * const x) {

        // put the elements of the antihermitian matrix's lower triangle in a vector
        _lt.reset(_nlt) <<= x;

        // build the anti-hermitian matrix from the lower triangle
        _ahm.reset(_norb, _norb, formic::zero(S()));
        for (int i = 0; i < _norb; i++)
        for (int j = 0; j < i; j++) {
          const S x = _lt.at(formic::cmpd_pair_index(j, i, _norb));
          _ahm(i,j) = x;
          _ahm(j,i) = -formic::conj(x);
        }

        //formic::of << "_ahm:" << std::endl;
        //for (int i = 0; i < _norb; i++) {
        //  for (int j = 0; j < _norb; j++)
        //    formic::of << boost::format("  %20.12f") % _ahm(i,j);
        //  formic::of << std::endl;
        //}
        //formic::of << std::endl;

        // exponentiate the anti-hermitian matrix to get the unitary orbital transformation
        _uot = formic::matrix_exponent(_ahm, 1.0e-15);

        //formic::of << "_uot:" << std::endl;
        //for (int i = 0; i < _norb; i++) {
        //  for (int j = 0; j < _norb; j++)
        //    formic::of << boost::format("  %20.12f") % _uot(i,j);
        //  formic::of << std::endl;
        //}
        //formic::of << std::endl;

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  transforms the integrals using the stored lower triangle variables
      ///
      ///  original one-electron integrals should already be stored in _oei_pq
      ///  original two-electron integrals should already be stored in _tei_pqrs
      /// \param[in]      n      number of parameters that specify the orbital transformation
      /// \param[in]      x      the elements of the antihermitian matrix's lower triangle, size _nlt
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void forward_transform(const int n, const double * const x) {

        // build the unitary orbital transformation
        this->build_uot(n, x);

        // transform 1st index of one-electron integrals
        _oei_qi = ( _uot * _oei_pq ).t();

        // transform 2nd index of one-electron integrals
        _oei_ij = ( _uot * _oei_qi ).t();

        // transform 1st index of two-electron integrals
        _tei_qrsi.reset(_norb, _norb*_norb*_norb);
        formic::tei_rotate_indices_left(_norb, &( _uot * _tei_pqrs ).at(0,0), &_tei_qrsi.at(0,0));

        // transform 2nd index of two-electron integrals
        _tei_rsij.reset(_norb, _norb*_norb*_norb);
        formic::tei_rotate_indices_left(_norb, &( _uot * _tei_qrsi ).at(0,0), &_tei_rsij.at(0,0));

        // transform 3rd index of two-electron integrals
        _tei_sijk.reset(_norb, _norb*_norb*_norb);
        formic::tei_rotate_indices_left(_norb, &( _uot * _tei_rsij ).at(0,0), &_tei_sijk.at(0,0));

        // transform 4th index of two-electron integrals
        _tei_ijkl.reset(_norb, _norb*_norb*_norb);
        formic::tei_rotate_indices_left(_norb, &( _uot * _tei_sijk ).at(0,0), &_tei_ijkl.at(0,0));

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  function to work backwards through the integral transformation, computing the
      ///         adjoint derivatives for the lower triangle variables
      ///
      ///  The forward_transform function should already have been run to generate its intermediates.
      ///  Additionally:
      ///    the wave function derivatives' contributions should already be in _uot_adj
      ///    the 2-index-transformed oei derivative adjoint should already be stored in _oei_adj_ij
      ///    the 4-index-transformed tei derivative adjoint should already be stored in _tei_adj_ijkl
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void reverse_transform() {

        if (false) {

          for (int i = 0; i < _norb; i++)
          for (int j = 0; j < _norb; j++)
          for (int p = 0; p < _norb; p++)
            _uot_adj(i,p) += 2.0 * _oei_adj_ij(i,j) * _oei_qi(p,j);

        } else {

          // compute 1-index-transformed oei derivative adjoint
          _oei_adj_qi = ( _oei_adj_ij * _uot ).t();

          // get unitary orbital transform derivative adjoint contribution from 2-index-transformed oei derivative adjoint
          _uot_adj += ( _oei_qi * _oei_adj_ij ).t();

          // get unitary orbital transform derivative adjoint contribution from 1-index-transformed oei derivative adjoint
          _uot_adj += ( _oei_pq * _oei_adj_qi ).t();

        }

        // compute 3-index-transformed tei derivative adjoint
        _tei_adj_sijk.reset(_norb*_norb*_norb, _norb);
        formic::tei_rotate_indices_right(_norb, &( _tei_adj_ijkl * _uot ).at(0,0), &_tei_adj_sijk.at(0,0));

        // compute 2-index-transformed tei derivative adjoint
        _tei_adj_rsij.reset(_norb*_norb*_norb, _norb);
        formic::tei_rotate_indices_right(_norb, &( _tei_adj_sijk * _uot ).at(0,0), &_tei_adj_rsij.at(0,0));

        // compute 1-index-transformed tei derivative adjoint
        _tei_adj_qrsi.reset(_norb*_norb*_norb, _norb);
        formic::tei_rotate_indices_right(_norb, &( _tei_adj_rsij * _uot ).at(0,0), &_tei_adj_qrsi.at(0,0));

        if (false) {

          for (int i = 0; i < _norb; i++)
          for (int j = 0; j < _norb; j++)
          for (int k = 0; k < _norb; k++)
          for (int l = 0; l < _norb; l++)
          for (int p = 0; p < _norb; p++)
            _uot_adj(i,p) += 4.0 * (    &_tei_sijk.at(0,0))[ p + _norb * k + _norb*_norb * l + _norb*_norb*_norb * j ]
                                 * (&_tei_adj_ijkl.at(0,0))[ i + _norb * j + _norb*_norb * k + _norb*_norb*_norb * l ];

        } else {

          // get unitary orbital transform derivative adjoint contribution from 4-index-transformed tei derivative adjoint
          _uot_adj += ( _tei_sijk * _tei_adj_ijkl ).t();
          //for (int i = 0; i < _norb; i++)
          //for (int j = 0; j < _norb; j++)
          //for (int k = 0; k < _norb; k++)
          //for (int l = 0; l < _norb; l++)
          //for (int s = 0; s < _norb; s++)
          //  _uot_adj(l,s) +=   (    &_tei_sijk.at(0,0))[ s + _norb * i + _norb*_norb * j + _norb*_norb*_norb * k ]
          //                   * (&_tei_adj_ijkl.at(0,0))[ i + _norb * j + _norb*_norb * k + _norb*_norb*_norb * l ];

          // get unitary orbital transform derivative adjoint contribution from 3-index-transformed tei derivative adjoint
          _uot_adj += ( _tei_rsij * _tei_adj_sijk ).t();
          //for (int i = 0; i < _norb; i++)
          //for (int j = 0; j < _norb; j++)
          //for (int k = 0; k < _norb; k++)
          //for (int r = 0; r < _norb; r++)
          //for (int s = 0; s < _norb; s++)
          //  _uot_adj(k,r) +=   (    &_tei_rsij.at(0,0))[ r + _norb * s + _norb*_norb * i + _norb*_norb*_norb * j ]
          //                   * (&_tei_adj_sijk.at(0,0))[ s + _norb * i + _norb*_norb * j + _norb*_norb*_norb * k ];

          // get unitary orbital transform derivative adjoint contribution from 2-index-transformed tei derivative adjoint
          _uot_adj += ( _tei_qrsi * _tei_adj_rsij ).t();
          //for (int i = 0; i < _norb; i++)
          //for (int j = 0; j < _norb; j++)
          //for (int q = 0; q < _norb; q++)
          //for (int r = 0; r < _norb; r++)
          //for (int s = 0; s < _norb; s++)
          //  _uot_adj(j,q) +=   (    &_tei_qrsi.at(0,0))[ q + _norb * r + _norb*_norb * s + _norb*_norb*_norb * i ]
          //                   * (&_tei_adj_rsij.at(0,0))[ r + _norb * s + _norb*_norb * i + _norb*_norb*_norb * j ];

          // get unitary orbital transform derivative adjoint contribution from 1-index-transformed tei derivative adjoint
          _uot_adj += ( _tei_pqrs * _tei_adj_qrsi ).t();
          //for (int i = 0; i < _norb; i++)
          //for (int p = 0; p < _norb; p++)
          //for (int q = 0; q < _norb; q++)
          //for (int r = 0; r < _norb; r++)
          //for (int s = 0; s < _norb; s++)
          //  _uot_adj(i,p) +=   (    &_tei_pqrs.at(0,0))[ p + _norb * q + _norb*_norb * r + _norb*_norb*_norb * s ]
          //                   * (&_tei_adj_qrsi.at(0,0))[ q + _norb * r + _norb*_norb * s + _norb*_norb*_norb * i ];

        }

        //formic::of << "_uot_adj:" << std::endl;
        //for (int i = 0; i < _norb; i++) {
        //  for (int p = 0; p < _norb; p++)
        //    formic::of << boost::format("  %20.12f") % _uot_adj(i,p);
        //  formic::of << std::endl;
        //}
        //formic::of << std::endl;

        // compute the derivative adjoint of the anti-hermitian matrix that was exponentiated
        _ahm_adj = formic::matrix_exponent_der_adj(_ahm, _uot_adj, 1.0e-15);

        // build the anti-hermitian matrix from the lower triangle
        _lt_adj.reset(_nlt, formic::zero(S()));
        for (int i = 0; i < _norb; i++)
        for (int j = 0; j < i; j++) {
          //const S x = _lt.at(formic::cmpd_pair_index(j, i, _norb));
          //_ahm(i,j) = x;
          //_ahm(j,i) = -formic::conj(x);
          S & x = _lt_adj.at(formic::cmpd_pair_index(j, i, _norb));
          x += _ahm_adj(i,j);
          x -= _ahm_adj(j,i);
        }

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  prepare to compute the energy
      ///
      /// \param[in]      n      number of parameters that specify the orbital transformation
      /// \param[in]      x      the value of the orbital transformation parameters
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void energy_prep(const int n, const double * const x) {

        // root process transforms the integrals
        if ( formic::mpi::rank() == 0 )
          this->forward_transform(n, x);

        // set the hamiltonian to use the transformed integrals
        if ( formic::mpi::rank() == 0 )
          _ham.set_integrals(&_oei_ij.at(0,0), &_tei_ijkl.at(0,0));
        else
          _ham.set_integrals(0,0);
        //_ham.from_archive("fqmc_orb_opt_init_ham.bin");
        //_ham.transform_orbitals(&_uot.at(0,0));

        //// update the wave function with the unitary orbital transformation
        //_co->from_vector(&_saved_wfn_vars.at(0));
        //_co->apply_orbital_rotation(_norb, &_uot.at(0,0));

        if (_couple_bfgs) {
          formic::Vector<S,int> wfn_params(_co->n_der());
          if ( formic::mpi::rank() == 0 )
            wfn_params <<= ( x + _nlt );
          formic::mpi::bcast(wfn_params);
          _co->from_vector(&wfn_params.at(0));
        }

      }

    // public member functions
    public:

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  constructs the orbital optimization worker
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      OrbOptWorker(const formic::InputBase & userinp,
                   const int burn_len,
                   const int samp_len,
                   const formic::Configuration & init_config,
                   const formic::LaggedFibonacci<> & rng,
                   const boost::shared_ptr<formic::CoeffObj<S> > & co,
                   formic::Hamiltonian<S> & ham)
        : _norb(userinp.get<int>("nao")),
          _nlt(formic::binom_coeff(userinp.get<int>("nao"), 2)),
          _couple_bfgs(userinp.get<bool>("oo_couple_bfgs")),
          _interstitial_wfn(userinp.get<bool>("oo_interstitial_wfn")),
          _rotate_fixed_pm(userinp.get<bool>("oo_rot_fixed_pm")),
          _new_sample_each_iter(userinp.get<bool>("oo_new_samp_each_iter")),
          _co(co),
          _old_co( _co->clone() ),
          _ham(ham),
          _userinp(userinp)
      {

        
        boost::shared_ptr<formic::fqmc::Sampler<S> > temp_sampler( new formic::fqmc::Sampler<S>(userinp, burn_len, samp_len, init_config, rng, _old_co) );
        _sampler = temp_sampler;

        if ( _interstitial_wfn && !userinp.get<bool>("exact_sampling") )
          throw formic::Exception("interstitial wave function optimization is currently only implemented for exact sampling");

        // root process saves the initial integrals
        if ( formic::mpi::rank() == 0 ) {
          _oei_pq.reset(_norb, _norb) <<= _ham.dense_oei_ptr();
          _tei_pqrs.reset(_norb, _norb*_norb*_norb) <<= _ham.dense_tei_ptr();
          _ham.to_archive("fqmc_orb_opt_init_ham.bin");
        }

        // if necessary, root process saves the initial pairing matrix
        if ( formic::mpi::rank() == 0 && _rotate_fixed_pm ) {

          // get all the variables in the coefficient object
          formic::Vector<S,int> initial_vars(_co->n_var());
          _co->all_to_vector(&initial_vars.at(0));

          // extract the pairing matrix, which should be the first data in the variable vector
          _initial_pm.reset(_norb, _norb);
          _initial_pm <<= &initial_vars.at(0);

        }

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  compute the energy for the supplied orbital transformation parameters
      ///
      /// \param[in]      n      number of parameters that specify the orbital transformation
      /// \param[in]      x      the elements of the antihermitian matrix's lower triangle
      /// \param[out]     val    on exit, the wave function's energy for the supplied transformation
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void get_value(const int n, const double * const x, double & val) {

        // load best parameter set
        if ( !_couple_bfgs ) {
          if ( formic::mpi::rank() == 0 && _pmt_set_history.size() > 0 ) {
            formic::Vector<double,int> x_vec = ( formic::Vector<double,int>(n) <<= x );
            formic::Vector<double,int> dist_vec = _x_history.at(0) - x_vec;
            double min_distance = ( dist_vec.t() * dist_vec ).at(0);
            _co->from_vector(&_pmt_set_history.at(0).at(0));
            for (int i = 1; i < _pmt_set_history.size(); i++) {
              dist_vec = _x_history.at(i) - x_vec;
              double dist = ( dist_vec.t() * dist_vec ).at(0);
              if ( dist < min_distance ) {
                min_distance = dist;
                _co->from_vector(&_pmt_set_history.at(i).at(0));
              }
            }
          }
          formic::Vector<S,int> wfn_params(_co->n_der(), formic::zero(S()));
          if ( formic::mpi::rank() == 0 )
            _co->to_vector(&wfn_params.at(0));
          if ( formic::mpi::rank() == 0 )
            formic::of << "pre_bc" << std::endl;
          formic::mpi::bcast(wfn_params);
          if ( formic::mpi::rank() == 0 )
            formic::of << "post_bc" << std::endl;
          _co->from_vector(&wfn_params.at(0));
        }

        // prepare for energy calculation
        this->energy_prep(n, x);

        // if requested, rotate the fixed pairing matrix into the current orbital basis
        if ( _rotate_fixed_pm ) {

          // check that rotation is allowed
          if ( _couple_bfgs )
            throw formic::Exception("rotating fixed pairing matrix in orbital optimization is not currently compatible with coupling the orbital and wave function bfgs");
          if ( _interstitial_wfn )
            throw formic::Exception("rotating fixed pairing matrix in orbital optimization is not currently compatible with interstitial wave function optimization");

          // get a vector of the current wave function variables
          formic::Vector<S,int> all_pmts(_co->n_var());
          _co->all_to_vector(&all_pmts.at(0));

          // root process performs the rotation and places the rotated pairing matrix in the front section of the variable vector
          if ( formic::mpi::rank() == 0 )
            _uot * _initial_pm * _uot.t() >>= &all_pmts.at(0);
            //_uot.t() * _initial_pm * _uot >>= &all_pmts.at(0);

          // all processes update the coefficient object's variables with the newly rotated pairing matrix
          formic::mpi::bcast(all_pmts);
          _co->all_from_vector(&all_pmts.at(0));

          //if ( formic::mpi::rank() == 0 ) {
          //  formic::of << "printing variables after rotating fixed pairing matrix" << std::endl;
          //  _co->print_variables();
          //  formic::of << std::endl;
          //}

        }

        // if we are not coupling the wfn and rotation within bfgs, then the wave function parameters are optimized for each rotation
        if ( !_couple_bfgs )
          formic::fqmc::optimize(_userinp, _sampler, _co, false);

        // construct the accumulator
        boost::shared_ptr<formic::Accumulator<S> > eacc( new formic::fqmc::EnAcc<S>(_old_co,
                                                                                    _co,
                                                                                    0,
                                                                                    _userinp.get<bool>("exact_sampling"),
                                                                                    _userinp.get<bool>("print_config_info")) );

        // sample the energy
        _sampler->simple_sample() << eacc << _co;

        // return the energy
        val = ( dynamic_cast<formic::fqmc::EnAcc<S> *>(eacc.get()) )->energy();

        // save the wfn parameters we found
        if ( !_couple_bfgs && formic::mpi::rank() == 0 ) {
          formic::Vector<S,int> pmts(_co->n_der());
          _co->to_vector(&pmts.at(0));
          _pmt_set_history.push_back(pmts);
          _x_history.push_back(formic::Vector<double,int>(n) <<= x);
        }

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  compute the gradient with respect to the antihermitian matrix's lower triangle
      ///
      /// \param[in]      n      number of parameters that specify the orbital transformation
      /// \param[in]      x      the elements of the antihermitian matrix's lower triangle
      /// \param[out]     g      on exit, the gradient
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void get_gradient(const int n, const double * const x, double * const g) {

        // prepare for adjoint derivatives calculation
        this->energy_prep(n, x);

        // construct the orbital optimization accumulator
        boost::shared_ptr<formic::Accumulator<S> > oo_acc( new formic::fqmc::OrbOptAcc<S>(_norb, _uot, _old_co, _co) );
        const formic::fqmc::OrbOptAcc<S> * const oo_acc_p = dynamic_cast<formic::fqmc::OrbOptAcc<S> *>(oo_acc.get());

        // construct the gradient accumulator
        boost::shared_ptr<formic::Accumulator<S> > gr_acc;
        const formic::fqmc::GradAcc<S> * gr_acc_p;
        if ( _couple_bfgs && _rotate_fixed_pm )
          throw formic::Exception("rotating fixed pairing matrix in orbital optimization is not currently compatible with coupling the orbital and wave function bfgs");
        if (_couple_bfgs) {
          boost::shared_ptr<formic::Accumulator<S> > temp_gr( new formic::fqmc::GradAcc<S>(_old_co, _co) );
          gr_acc = temp_gr;
          gr_acc_p = dynamic_cast<formic::fqmc::GradAcc<S> *>(gr_acc.get());
        }
        if (_rotate_fixed_pm) {
          boost::shared_ptr<formic::Accumulator<S> > temp_gr( new formic::fqmc::GradAcc<S>(_old_co, _co, false) );
          gr_acc = temp_gr;
          gr_acc_p = dynamic_cast<formic::fqmc::GradAcc<S> *>(gr_acc.get());
        }

        // sample the derivatives we need
        if ( _couple_bfgs || _rotate_fixed_pm )
          _sampler->simple_sample() << oo_acc << gr_acc << _co;
        else
          _sampler->simple_sample() << oo_acc << _co;

        // root process computes the gradient via reverse-transformation
        if ( formic::mpi::rank() == 0 ) {

          // get the one-electron integral derivative adjoints
          _oei_adj_ij <<= oo_acc_p->oei_der();

          // get the two-electron integral derivative adjoints
          formic::tei_compact_to_dense(_norb, &oo_acc_p->tei_cmp_der().at(0), &_tei_adj_ijkl.reset(_norb*_norb*_norb, _norb).at(0,0));

          //for (int i = 0; i < _norb; i++)
          //for (int j = 0; j < _norb; j++)
          //  _oei_adj_ij(i,j) *= ( i == j ? 1.0 : 0.5 );

          const bool print_derivatives = false;

          if (print_derivatives) {

            formic::of << "_oei_adj_ij:" << std::endl;
            for (int i = 0; i < _norb; i++) {
              for (int j = 0; j < _norb; j++)
                formic::of << boost::format("  %20.12f") % _oei_adj_ij(i,j);
              formic::of << std::endl;
            }
            formic::of << std::endl;

            const S * const to_print = &_tei_adj_ijkl.at(0,0);
            formic::of << "_tei_adj_ijkl:" << std::endl;
            for (int l = 0; l < _norb; l++)
            for (int k = 0; k < _norb; k++)
            for (int j = 0; j < _norb; j++)
            for (int i = 0; i < _norb; i++)
              formic::of << boost::format("  %2i  %2i  %2i  %2i  %20.12f") % i % j % k % l % to_print[ i + _norb*j + _norb*_norb*k + _norb*_norb*_norb*l ] << std::endl;
            formic::of << std::endl;

          }

          // initialize the derivative adjoint for the unitary orbital transformation
          _uot_adj.reset(_norb, _norb, formic::zero(S()));

          // when rotating a fixed pairing matrix, get its contribution to the unitary orbital transformation's derivative adjoint
          if ( _rotate_fixed_pm ) {

            // get the gradient w.r.t. wave function variables
            formic::Vector<double,int> wfn_real_grad(2*_co->n_var());
            gr_acc_p->get_real_gradient(&wfn_real_grad.at(0));

            // convert to S in order to compile (note this is not a correct handling for complex numbers!)
            formic::Vector<S,int> wfn_grad(_co->n_var());
            for (int i = 0; i < wfn_grad.size(); i++)
              wfn_grad[i] = formic::unity(S()) * wfn_real_grad[i];

            // get the gradient w.r.t. the pairing matrix
            formic::ConstMatrix<S,int> pm_grad(_norb, _norb, &wfn_grad.at(0));

            // expression used to build rotated pairing matrix:   _uot * _initial_pm * _uot.t()

            // get the contribution to the unitary transform's derivative adjoint
            _uot_adj += pm_grad * _uot * _initial_pm.t();
            _uot_adj += pm_grad.t() * _uot * _initial_pm;

          }

          // compute the energy derivatives w.r.t. the orbital rotation parameters
          this->reverse_transform();
          
          //_lt_adj *= -1.0;

          // return the energy derivatives w.r.t. the orbital rotation parameters
          _lt_adj >>= g;

          // return the energy derivatives w.r.t. wave function parameters
          if (_couple_bfgs)
            gr_acc_p->get_real_gradient(g + _nlt);

        }

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  print a set of lower triangle variables
      ///
      /// \param[in]      x      the elements of the antihermitian matrix's lower triangle, size _nlt
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void print_lower_triangle(const S * const x) {
        for (int i = 0; i < _norb; i++) {
          for (int j = 0; j < i; j++)
            formic::of << boost::format("  %12.6f") % x[formic::cmpd_pair_index(j, i, _norb)];
          formic::of << boost::format("  %12.6f") % formic::zero(S()) << std::endl;
        }
        formic::of << std::endl;
      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  print function used by lbfgs
      ///
      /// \param[in]      n        number of parameters that specify the orbital transformation
      /// \param[in]      message  a description of what is being printed
      /// \param[in]      x        the elements of the antihermitian matrix's lower triangle, size _nlt
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void print(const int n, const std::string & message, const double * const x) {
        //if ( formic::mpi::rank() == 0 ) {
        //  formic::of << message << std::endl;
        //  formic::of << std::endl;
        //  formic::of << "lower triangle of antihermitian matrix:" << std::endl;
        //  this->print_lower_triangle(x);
        //  if (_couple_bfgs) {
        //    _co->from_vector(x + _nlt);
        //    _co->print_variables();
        //  }
        //}
      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  if optimizes the wave function variables at the new orbital point
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void process_new_point(const int n, double * const x) {

        // print new orbital basis
        const bool from_curr = true;
        const bool print_only = true;
        if ( formic::mpi::rank() == 0 ) {
          this->build_uot(n, x);
          formic::of << boost::format("new orbital basis is:") << std::endl;
          _ham.transform_orbitals(&_uot.at(0,0), from_curr, print_only);
        } else {
          _ham.transform_orbitals(0, from_curr, print_only);
        }

        // if requested, update the guiding wave function and take a new MC sample for the next iteration
        if ( _new_sample_each_iter ) {
          _old_co = _co->clone();
          _sampler->set_gf(_old_co);
        }

        // reset the history of known parameter sets
        _pmt_set_history.clear();
        _x_history.clear();

        // if requested, optimize the wave function parameters for the new integrals
        if (_interstitial_wfn) {
          formic::Vector<S,int> wfn_params(_co->n_der());
          if (_couple_bfgs) {
            if ( formic::mpi::rank() == 0 )
              wfn_params <<= x+_nlt;
            formic::mpi::bcast(wfn_params);
            _co->from_vector(&wfn_params.at(0));
          }
          formic::fqmc::optimize(_userinp, _sampler, _co, false);
          if ( _couple_bfgs && formic::mpi::rank() == 0 ) {
            _co->to_vector(&wfn_params.at(0));
            wfn_params >>= x+_nlt;
          }
        }

      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief  transforms the hamiltonian with the specified unotary orbital rotation
      ///
      /// \param[in]      n        number of parameters that specify the orbital transformation
      /// \param[in]      x        the elements of the antihermitian matrix's lower triangle, size _nlt
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      void transform_hamiltonian(const int n, const double * const x) {

        // root process computes the unitary orbital transformation 
        if ( formic::mpi::rank() == 0 )
          this->build_uot(n, x);

        //// load the unadaltered hamiltonian
        //_ham.from_archive("fqmc_orb_opt_init_ham.bin");

        // transformt the hamiltonian
        if ( formic::mpi::rank() == 0 )
          _ham.transform_orbitals(&_uot.at(0,0));
        else
          _ham.transform_orbitals(0);

      }

  };

} // end namespace fqmc

} // end namespace formic

template<class W, class S> static void my_fd_test(W & worker,
                                                  const double iv,
                                                  formic::Vector<S,int> & x,
                                                  formic::Vector<S,int> & g,
                                                  const int ind,
                                                  const std::string & name) {

  std::vector<double> dxv;
  dxv.push_back(1.0e-4);
  dxv.push_back(1.0e-5);
  dxv.push_back(1.0e-6);
  dxv.push_back(1.0e-7);
  dxv.push_back(1.0e-8);
  dxv.push_back(1.0e-9);

  std::vector<double> fdv;
  for (int i = 0; i < dxv.size(); i++) {
    const S xold = x.at(ind);
    x.at(ind) += dxv.at(i) * formic::unity(S());
    double nv;
    worker.get_value(x.size(), &x.at(0), nv);
    x.at(ind) = xold;
    fdv.push_back( ( nv - iv ) / dxv.at(i) );
  }

  formic::of << boost::format("%40s                  ") % " ";
  for (int i = 0; i < dxv.size(); i++)
    formic::of << boost::format(" %20.2e") % dxv.at(i);
  formic::of << std::endl;

  formic::of << boost::format("%40s analytic der    =") % name;
  for (int i = 0; i < dxv.size(); i++)
    formic::of << boost::format(" %20.12f") % formic::real(g.at(ind));
  formic::of << std::endl;

  formic::of << boost::format("%40s finite diff der =") % name;
  for (int i = 0; i < dxv.size(); i++)
    formic::of << boost::format(" %20.12f") % fdv.at(i);
  formic::of << std::endl;

  formic::of << boost::format("%40s relative error  =") % name;
  for (int i = 0; i < dxv.size(); i++)
    formic::of << boost::format(" %20.2e") % ( ( formic::real(g.at(ind)) - fdv.at(i) ) / fdv.at(i) );
  formic::of << std::endl;

  formic::of << std::endl;

}

template<class S> void formic::fqmc::optimize_orbitals(const formic::InputBase & userinp,
                                                       formic::Hamiltonian<S> & ham,
                                                       const boost::shared_ptr<formic::CoeffObj<S> > & co) {

  // create a worker for the lbfgs optimization
  formic::fqmc::OrbOptWorker<S> worker(userinp,
                                       userinp.get<int>("opt_burn_len"),
                                       userinp.get<int>("opt_samp_len"), 
                                       formic::Configuration(userinp),
                                       formic::LaggedFibonacci<>(),
                                       co,
                                       ham);

  const bool couple_bfgs = userinp.get<bool>("oo_couple_bfgs");

  const int norb = userinp.get<int>("nao");
  const int nlt  = formic::binom_coeff(norb, 2);
  const int nvar = nlt + ( couple_bfgs ? co->n_der() : 0 );

  formic::Vector<S,int> x0(nvar, formic::zero(S()));
  formic::Vector<S,int> x1(nvar, formic::zero(S()));

  //for (int i = 0; i < x0.size(); i++)
  //  x0.at(i) = 0.10 * 2.0 * ( formic::random_number<S>() - 0.5 * formic::unity(S()) );

  if (couple_bfgs)
    co->to_vector(&x0.at(nlt));

  // optimize the orbital rotation
  formic::lbfgs(worker,
                userinp.get<int>("orb_opt_max_iter"),
                //nvar*8/9,
                1000,
                nvar,
                userinp.get<double>("oo_gr_thresh"),
                userinp.get<double>("oo_en_thresh"),
                &x0.at(0),
                &x1.at(0));

  //// optimize the orbital rotation
  //formic::simple_sr1(worker,
  //                   userinp.get<int>("orb_opt_max_iter"),
  //                   userinp.get<double>("opt_gr_thresh"),
  //                   x0,
  //                   x1);

  // check finite difference
  static bool check_fd = false;
  if ( check_fd ) {

    formic::of << "!!! Finite Difference Check !!!" << std::endl;

    double initial_val;
    worker.get_value(nvar, &x1.at(0), initial_val);
    worker.get_gradient(nvar, &x1.at(0), &x0.at(0));

    //const S dx = 1.0e-4 * formic::unity(S());

    for (int i = 0; i < norb; i++)
    for (int j = 0; j < i; j++)
      my_fd_test(worker, initial_val, x1, x0, formic::cmpd_pair_index(j, i, norb), (boost::format("orbital rotation %2i %2i") % i % j).str());
    formic::of << std::endl;

    for (int i = 0; i < norb; i++)
    for (int j = 0; j < norb; j++)
      my_fd_test(worker, initial_val, x1, x0, nlt+i*norb+j, (boost::format("pm element %2i %2i") % i % j).str());
    formic::of << std::endl;

    for (int i = nlt+norb*norb; i < nvar; i++)
      my_fd_test(worker, initial_val, x1, x0, i, (boost::format("jf element %2i") % (i-nlt-norb*norb)).str());

  }
  check_fd = false;

  // apply the optimized orbital rotation
  worker.transform_hamiltonian(nlt, &x1.at(0));

  // get the wave function parameters
  if (couple_bfgs) {
    formic::mpi::bcast(x1);
    co->from_vector(&x1.at(nlt));
  }

}

template void formic::fqmc::optimize_orbitals(const formic::InputBase & userinp,
                                              formic::Hamiltonian<double> & ham,
                                              const boost::shared_ptr<formic::CoeffObj<double> > & co);
//template void formic::fqmc::optimize_orbitals(const formic::InputBase & userinp,
//                                              formic::Hamiltonian<std::complex<double> > & ham,
//                                              const boost::shared_ptr<formic::CoeffObj<std::complex<double> > > & co);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "OptimizeOrbitals" command, which minimizes the energy
///          by changing the one electron orbital basis.
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void formic::fqmc::OptimizeOrbitals_command_function(formic::Archive & arch, formic::AllData & data) {

  // read name of coefficient object to use
  std::string co_name;
  formic::mpi::read_and_bcast(arch, co_name, "Failed to read co_name in formic::fqmc::OptimizeOrbitals_command_function.");

  // execute the optimization
  if (data.userinp.get<std::string>("scalar_type") == "real") {
    formic::fqmc::optimize_orbitals(data.userinp, data.hd, data.wd.get_co(co_name));
  } else if (data.userinp.get<std::string>("scalar_type") == "complex") {
    throw formic::Exception("command \"OptimizeOrbitals\" is currently only implemented for real numbers");
  } else
    throw formic::Exception("unknown scalar type \"%s\" when executing command \"OptimizeOrbitals\"") % data.userinp.get<std::string>("scalar_type");

}
